package com.example.emqx.config;

import jakarta.servlet.WriteListener;
import jakarta.servlet.http.HttpServletResponse;
import jakarta.servlet.http.HttpServletResponseWrapper;

import java.io.ByteArrayOutputStream;
import java.io.CharArrayWriter;
import java.io.IOException;
import java.io.PrintWriter;
import jakarta.servlet.ServletOutputStream;

/**
 * 自定义 HttpServletResponseWrapper 类，用于捕获响应数据
 */
public class CharResponseWrapper extends HttpServletResponseWrapper {
    private final CharArrayWriter charWriter = new CharArrayWriter();
    private PrintWriter writer;
    private ServletOutputStream outputStream;

    public CharResponseWrapper(HttpServletResponse response) {
        super(response);
    }

    @Override
    public PrintWriter getWriter() {
        if (writer == null) {
            writer = new PrintWriter(charWriter);
        }
        return writer;
    }

    @Override
    public ServletOutputStream getOutputStream() throws IOException {
        if (outputStream == null) {
            outputStream = new ServletOutputStream() {
                @Override
                public boolean isReady() {
                    return true;
                }

                @Override
                public void setWriteListener(WriteListener writeListener) {

                }

                private final ByteArrayOutputStream output = new ByteArrayOutputStream();

                @Override
                public void write(int b) throws IOException {
                    output.write(b);
                }

                @Override
                public void flush() throws IOException {
                    output.flush();
                }

                @Override
                public void close() throws IOException {
                    output.close();
                }

                public String toString() {
                    return output.toString();
                }
            };
        }
        return outputStream;
    }

    @Override
    public String toString() {
        if (outputStream != null) {
            return outputStream.toString();  // 通过 OutputStream 捕获的数据
        } else {
            return charWriter.toString();  // 通过 Writer 捕获的数据
        }
    }

    public void writeResponse() throws IOException {
        HttpServletResponse response = (HttpServletResponse) getResponse();
        response.getWriter().write(toString());
    }
}
